package main

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"mime"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

const URL = "https://download-alibaba.jetbrains.com.cn/go/goland-2024.3.1.exe"
const PART_NUM = 5 // 下载的线程数

type Downloader struct {
	url       string
	outputdir string
	filename  string
	filesize  int
}

func NewDownloader(urlStr, outputdir, filename string) *Downloader {
	return &Downloader{
		url:       urlStr,
		outputdir: outputdir,
		filename:  filename,
	}
}

func (d *Downloader) head() error {
	client := http.Client{}
	res, _ := http.NewRequest(http.MethodHead, d.url, nil)
	// res.ParseForm()
	resp, err := client.Do(res)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.Header.Get("Accept-Ranges") != "bytes" {
		return fmt.Errorf("服务器不支持断点续传")
	}

	filesize := resp.Header.Get("Content-Length")
	if filesize == "" {
		return fmt.Errorf("文件错误")
	}
	d.filesize, _ = strconv.Atoi(filesize)

	disposition := resp.Header.Get("Content-Disposition")
	if disposition != "" {
		_, params, _ := mime.ParseMediaType(disposition)
		d.filename = params["filename"]
	} else {
		d.filename = filepath.Base(d.url)
	}

	d.outputdir, _ = os.Getwd()
	return nil
}

func (d *Downloader) run() error {
	err := d.head()
	if err != nil {
		return err
	}

	fmt.Println("开始下载：" + d.url)

	base := d.filesize / PART_NUM
	partFiles := make([]string, PART_NUM)

	var wg sync.WaitGroup
	for i := 0; i < PART_NUM; i++ {
		begin := 0
		if i != 0 {
			begin = i*base + 1
		}

		end := (i + 1) * base
		if i == PART_NUM-1 {
			end = d.filesize
		}

		partFilename := fmt.Sprintf("%s.part%d", d.filename, i)
		partFiles[i] = filepath.Join(d.outputdir, partFilename)

		wg.Add(1)
		go func(i, start, end int, partFilename string) {
			defer wg.Done()
			err := d.partDown(i, begin, end, partFilename)
			if err != nil {
				fmt.Printf("线程:%d 下载失败: %v\n", i+1, err)
			} else {
				fmt.Printf("线程:%d 下载成功\n", i+1)
			}
		}(i, begin, end, partFilename)
	}

	wg.Wait()

	fmt.Println("开始合成数据...")
	d.mergeParts(partFiles)
	fmt.Println("合并完成")

	for _, partFile := range partFiles {
		err := os.Remove(partFile)
		if err != nil {
			// 如果删除失败，打印错误信息
			fmt.Printf("Error removing directory: %v\n", err)
		}
	}
	return nil
}

func (d *Downloader) mergeParts(partFiles []string) error {
	out, err := os.Create(filepath.Join(d.outputdir, d.filename))
	if err != nil {
		return err
	}
	defer out.Close()

	for _, partFile := range partFiles {
		in, err := os.Open(partFile)
		if err != nil {
			return err
		}
		defer in.Close()

		_, err = io.Copy(out, in)
		if err != nil {
			return err
		}
	}

	return nil

	// hashString := hex.EncodeToString(hash.Sum(nil))
	// if hashString != "97e88c1e3c8d83a543aade719d2edf3c74994c698be0343b629e48306e4a532a" {
	// 	fmt.Println(hashString)
	// }

}

func (d *Downloader) partDown(i, start, end int, filename string) error {
	req, _ := http.NewRequest("GET", d.url, nil)
	// req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36")
	req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", start, end))

	fmt.Printf("线程:%d 开始下载>\r\n", (i + 1))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusPartialContent {
		return fmt.Errorf("unexpected status code: %d", resp.StatusCode)
	}

	out, err := os.Create(filename)
	if err != nil {
		return err
	}

	_, err = io.Copy(out, resp.Body)

	return err
}

func main() {
	fmt.Printf("请输入URL：")
	reader := bufio.NewReader(os.Stdin)
	var urlStr string
	for {
		urlStr, _ = reader.ReadString('\n')
		urlStr = strings.TrimSpace(urlStr)
		err := checkUrl(urlStr)
		if err != nil {
			fmt.Println(err)
		} else {
			break
		}
	}

	downloader := NewDownloader(urlStr, "", "")
	startTime := time.Now()
	downloader.run()
	duration := time.Since(startTime)
	fmt.Println()
	fmt.Printf("程序运行了 %v\n", duration)
}

func checkUrl(urlStr string) error {
	u1, err := url.ParseRequestURI(urlStr)
	if err != nil {
		return errors.New(err.Error() + "; 请重新输入URL：")
	}

	if u1.Host == "" || u1.Scheme == "" {
		return fmt.Errorf("url 格式错误" + "; 请重新输入URL：")
	}

	return nil
}
